import os
import pandas as pd
import pymysql
from pymysql.cursors import DictCursor
from datetime import datetime

# 数据库配置（请替换为你的实际配置）
DB_CONFIG = {
    "host": "localhost",
    "user": "root",
    "password": "你的数据库密码",
    "db": "stock_db",  # 数据库名
    "port": 3306,
    "charset": "utf8mb4"
}

# CSV文件路径与数据库表的映射关系
# 键：CSV文件所在文件夹路径；值：对应的数据库表名
CSV_TABLE_MAP = {
    "./daily_data": "stock_daily",  # 日线数据CSV文件夹 → 对应stock_daily表
    "./minute_data": "stock_minute",  # 分钟线数据CSV文件夹 → 对应stock_minute表
    "./basic_info": "stock_basic",  # 基础信息CSV文件夹 → 对应stock_basic表
    "./finance_data": "stock_finance"  # 财务数据CSV文件夹 → 对应stock_finance表
}

# 字段映射：CSV列名到数据库字段名的映射（如果不一致需要手动映射）
FIELD_MAPPING = {
    "stock_daily": {
        "股票代码": "stock_code",
        "日期": "trade_date",
        "开盘": "open_price",
        "收盘": "close_price",
        "最高": "high_price",
        "最低": "low_price",
        "成交量": "volume",
        "成交额": "amount",
        "涨跌幅": "change_rate",
        "换手率": "turnover_rate"
    },
    # 其他表的字段映射可以根据实际CSV列名添加
    "stock_basic": {
        "股票代码": "stock_code",
        "股票名称": "stock_name",
        "交易所": "exchange",
        "行业": "industry",
        "上市日期": "list_date",
        "总股本": "total_share",
        "流通股本": "circulating_share"
    }
}


def get_db_connection():
    """获取数据库连接"""
    try:
        conn = pymysql.connect(
            host=DB_CONFIG["host"],
            user=DB_CONFIG["user"],
            password=DB_CONFIG["password"],
            db=DB_CONFIG["db"],
            port=DB_CONFIG["port"],
            charset=DB_CONFIG["charset"],
            cursorclass=DictCursor
        )
        return conn
    except Exception as e:
        print(f"数据库连接失败：{str(e)}")
        return None


def get_table_columns(conn, table_name):
    """获取表的字段列表"""
    try:
        with conn.cursor() as cursor:
            cursor.execute(f"DESCRIBE {table_name}")
            columns = [row["Field"] for row in cursor.fetchall()]
        return columns
    except Exception as e:
        print(f"获取表{table_name}字段失败：{str(e)}")
        return []


def process_csv_data(df, table_name):
    """处理CSV数据，转换格式并映射字段名"""
    # 字段名映射
    if table_name in FIELD_MAPPING:
        df = df.rename(columns=FIELD_MAPPING[table_name])

    # 处理日期格式
    if "trade_date" in df.columns:
        df["trade_date"] = pd.to_datetime(df["trade_date"]).dt.date
    if "list_date" in df.columns:
        df["list_date"] = pd.to_datetime(df["list_date"], errors="coerce").dt.date

    # 处理分钟线时间
    if "trade_time" in df.columns:
        df["trade_time"] = pd.to_datetime(df["trade_time"], errors="coerce")

    # 填充缺失值
    df = df.fillna(0)

    return df


def import_csv_to_mysql(csv_file, table_name, conn):
    """将单个CSV文件导入到MySQL表"""
    try:
        # 读取CSV文件
        df = pd.read_csv(csv_file)
        print(f"读取文件 {csv_file}，共 {len(df)} 条数据")

        # 处理数据
        df = process_csv_data(df, table_name)

        # 获取表字段
        table_columns = get_table_columns(conn, table_name)
        if not table_columns:
            return False

        # 过滤出表中存在的字段
        df = df[[col for col in df.columns if col in table_columns]]
        if df.empty:
            print(f"没有匹配的字段，跳过文件 {csv_file}")
            return True

        # 构建插入SQL
        columns = ", ".join(df.columns)
        placeholders = ", ".join(["%s"] * len(df.columns))
        sql = f"""
            INSERT INTO {table_name} ({columns})
            VALUES ({placeholders})
            ON DUPLICATE KEY UPDATE
            {", ".join([f"{col}=VALUES({col})" for col in df.columns if col not in ["id", "created_time", "updated_time"]])}
        """

        # 执行批量插入
        data = [tuple(row) for _, row in df.iterrows()]
        with conn.cursor() as cursor:
            cursor.executemany(sql, data)
        conn.commit()
        print(f"成功导入 {len(data)} 条数据到表 {table_name}")
        return True

    except Exception as e:
        conn.rollback()
        print(f"导入文件 {csv_file} 失败：{str(e)}")
        return False


def batch_import():
    """批量导入所有CSV文件"""
    conn = get_db_connection()
    if not conn:
        return

    try:
        # 遍历所有配置的文件夹和表
        for csv_dir, table_name in CSV_TABLE_MAP.items():
            if not os.path.exists(csv_dir):
                print(f"文件夹 {csv_dir} 不存在，跳过")
                continue

            # 遍历文件夹中的所有CSV文件
            for filename in os.listdir(csv_dir):
                if filename.endswith(".csv"):
                    csv_path = os.path.join(csv_dir, filename)
                    print(f"\n开始处理 {csv_path}")
                    import_csv_to_mysql(csv_path, table_name, conn)

        print("\n所有文件处理完毕")

    finally:
        if conn:
            conn.close()


if __name__ == "__main__":
    # 执行批量导入
    batch_import()
